#include <iostream>
#include <Eigen/Dense>

using namespace Eigen;
using namespace std;

class LeastSquaresSolverQR {
private:
    VectorXd x;
    VectorXd y;
    MatrixXd A;
    MatrixXd Q;
    MatrixXd R;

public:
    LeastSquaresSolverQR(VectorXd& inputX, VectorXd& inputY) : x(inputX), y(inputY) {
        A = MatrixXd(inputX.size(), 3);
        A.col(0) = VectorXd::Ones(inputX.size());
        A.col(1) = inputX;
        A.col(2) = inputX.array().square();
    }

    VectorXd solve() {
        HouseholderQR<MatrixXd> qr;
        qr.compute(A);
        VectorXd coefficients = qr.solve(y);
        Q = qr.householderQ();
        R = qr.matrixQR().triangularView<Upper>();
        return coefficients;
    }

    double getConditionNumberG() const {
        MatrixXd G = A.transpose() * A;
        double conditionNumberG = G.norm() * G.inverse().norm();
        return conditionNumberG;
    }

    double getConditionNumberR1() const {
        MatrixXd R1 = R.block(0, 0, 3, 3);
        double conditionNumberR1 = R1.norm() * R1.inverse().norm();
        return conditionNumberR1;
    }
};

int main() {
    VectorXd xData(21), yData(21);
    xData << 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0;
    yData << 2.9, 2.7, 4.8, 5.3, 7.1, 7.6, 7.7, 7.6, 9.4, 9.0, 9.6, 10.0, 10.2, 9.7, 8.3, 8.4, 9.0, 8.3, 6.6, 6.7, 4.1;

    LeastSquaresSolverQR solver(xData, yData);
    VectorXd coefficients = solver.solve();
    cout << "Coefficients of the best-fit quadratic polynomial obtained by QR factorization:" << endl;
    cout << coefficients << endl;

    cout << "Condition number of the matrix G in 2-norm:" << endl;
    cout << solver.getConditionNumberG() << endl;

    cout << "Condition number of the matrix R1 in 2-norm:" << endl;
    cout << solver.getConditionNumberR1() << endl;

    return 0;
}
